#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import cv2
import numpy as np
import serial
import time
import threading
import math  
from rclpy.executors import MultiThreadedExecutor


# -------------------------- 协议与参数常量 --------------------------
FRAME_HEADER = 0x66
FRAME_LENGTH = 0x07
FRAME_TAIL = 0x5B
SEND_INTERVAL = 0.02  # 50Hz发送频率
CENTER_THRESHOLD = 2  # 中心阈值（像素）
MIN_CONTOUR_AREA = 100  # 最小轮廓面积（过滤小噪声）
MIN_CONTOUR_PERIMETER = 20  # 最小轮廓周长


# -------------------------- 双目测距相关函数 --------------------------
def init_camera_params():
    left_camera_matrix = np.array([[715.088662856028, -1.27750597113329, 321.725929157012],
                                   [0, 717.938092436680, 251.352481762661],
                                   [0., 0., 1.]])
    right_camera_matrix = np.array([[712.537452801351, -1.49459241948369, 325.232106968860],
                                    [0, 715.552769066380, 250.629910089625],
                                    [0., 0., 1.]])

    left_distortion = np.array(
        [[-0.440986151974959, 0.327464329751247, -0.000547150225973162, -0.00126940700042557, -0.335074894313998]])
    right_distortion = np.array(
        [[-0.444122749137210, 0.348225581098782, 0.00176792559350367, 0.00123611244757822, -0.660993977721134]])

    R = np.array([[0.999866011965132, 0.000416637643297468, 0.0163641232584165],
                  [-0.000544744807817963, 0.999969236927380, 0.00782486114692518],
                  [-0.0163603597159969, -0.00783272698033651, 0.999835479975588]])
    T = np.array([-59.8387197927031, 0.0107452720401373, 0.330883446596301])
    size = (640, 480)

    R1, R2, P1, P2, Q, validPixROI1, validPixROI2 = cv2.stereoRectify(
        left_camera_matrix, left_distortion,
        right_camera_matrix, right_distortion,
        size, R, T
    )

    left_map1, left_map2 = cv2.initUndistortRectifyMap(
        left_camera_matrix, left_distortion, R1, P1, size, cv2.CV_16SC2)
    right_map1, right_map2 = cv2.initUndistortRectifyMap(
        right_camera_matrix, right_distortion, R2, P2, size, cv2.CV_16SC2)

    return {
        "left_map1": left_map1, "left_map2": left_map2,
        "right_map1": right_map1, "right_map2": right_map2,
        "Q": Q, "size": size
    }


def init_stereo_matcher(blockSize=5, numDisparities=128):
    img_channels = 3
    stereo = cv2.StereoSGBM_create(
        minDisparity=0,
        numDisparities=numDisparities,
        blockSize=blockSize,
        P1=8 * img_channels * blockSize * blockSize,
        P2=32 * img_channels * blockSize * blockSize,
        disp12MaxDiff=5,
        preFilterCap=63,
        uniquenessRatio=15,
        speckleWindowSize=200,
        speckleRange=32,
        mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY
    )

    right_matcher = cv2.ximgproc.createRightMatcher(stereo)
    wls_filter = cv2.ximgproc.createDisparityWLSFilter(matcher_left=stereo)
    wls_filter.setLambda(8000)
    wls_filter.setSigmaColor(1.5)

    return stereo, right_matcher, wls_filter


def rectify_image(img, map1, map2):
    return cv2.remap(img, map1, map2, cv2.INTER_LINEAR)


def compute_disparity_and_depth(left_rect, right_rect, stereo, right_matcher, wls_filter, Q, depth_scale=16.0):
    disparity_left = stereo.compute(left_rect, right_rect)
    disparity_right = right_matcher.compute(right_rect, left_rect)

    filtered_disparity = wls_filter.filter(disparity_left, left_rect, None, disparity_right)
    filtered_disparity = np.clip(filtered_disparity, 0, None)

    disp_norm = cv2.normalize(filtered_disparity, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
    disp_color = cv2.applyColorMap(disp_norm, cv2.COLORMAP_JET)

    threeD = cv2.reprojectImageTo3D(filtered_disparity, Q, handleMissingValues=False)
    threeD = threeD * depth_scale

    return filtered_disparity, disp_norm, disp_color, threeD


# -------------------------- 主节点类 --------------------------
class BallDetectorNode(Node):
    def __init__(self):
        super().__init__('ball_detector_node')
        
        # 参数配置
        self.declare_parameter('camera_topic', '/image_raw')
        self.declare_parameter('serial_port', '/dev/ttyUSB0')
        self.declare_parameter('display_enabled', True)
        self.declare_parameter('center_threshold', CENTER_THRESHOLD)
        
        # 获取参数
        camera_topic = self.get_parameter('camera_topic').get_parameter_value().string_value
        serial_port = self.get_parameter('serial_port').get_parameter_value().string_value
        self.display_enabled = self.get_parameter('display_enabled').get_parameter_value().bool_value
        self.center_threshold = self.get_parameter('center_threshold').get_parameter_value().integer_value
        
        # 初始化组件
        self.bridge = CvBridge()
        self.image_sub = self.create_subscription(
            Image, camera_topic, self.image_callback, 10
        )
        self.ser = self.init_serial(serial_port)
        
        # 目标坐标和线程锁
        self.target_x = 320
        self.target_y = 240
        self.ball_distance = None
        self.coord_lock = threading.Lock()
        
        # 双目测距组件
        self.cam_params = init_camera_params()
        self.stereo, self.right_matcher, self.wls_filter = init_stereo_matcher()
        
        # 显示窗口
        if self.display_enabled:
            cv2.namedWindow("Ball Detection", cv2.WINDOW_NORMAL)
            cv2.namedWindow("Depth Map", cv2.WINDOW_NORMAL)
        
        # 串口发送线程
        self.send_thread = threading.Thread(target=self.send_coordinates_loop)
        self.send_thread.daemon = True
        self.send_thread.start()
        
        self.get_logger().info("Ball detector node started")

    def init_serial(self, port):
        try:
            ser = serial.Serial(
                port=port,
                baudrate=115200,
                bytesize=serial.EIGHTBITS,
                parity=serial.PARITY_NONE,
                stopbits=serial.STOPBITS_ONE,
                timeout=0.02
            )
            self.get_logger().info(f"Serial port {port} opened")
            return ser
        except serial.SerialException as e:
            self.get_logger().error(f"Serial error: {e}")
            return None

    # -------------------------- 优化的圆心检测函数（仅修改判断部分） --------------------------
    def detect_orange_ball(self, frame):
        """仅优化球体判断逻辑：强化轮廓筛选与圆心计算"""
        # 降采样加速
        small_frame = cv2.resize(frame, (320, 240))
        
        # 高斯模糊减少噪声
        blurred = cv2.GaussianBlur(small_frame, (5, 5), 0)
        hsv = cv2.cvtColor(blurred, cv2.COLOR_BGR2HSV)

        # 橙色阈值（保持用户设定）
        lower_orange = np.array([13, 64, 190])
        upper_orange = np.array([64, 255, 255])
        mask = cv2.inRange(hsv, lower_orange, upper_orange)

        # 形态学操作（保持原有参数）
        kernel = np.ones((5, 5), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=3)  # 填充空洞
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1)   # 去除噪点

        # 查找轮廓
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        best_score = 0
        ball_center = None

        for contour in contours:
            # 1. 面积判断：过滤过小轮廓
            area = cv2.contourArea(contour)
            if area < MIN_CONTOUR_AREA:
                continue  # 面积不达标直接跳过

            # 2. 周长判断：过滤过短轮廓
            perimeter = cv2.arcLength(contour, True)
            if perimeter < MIN_CONTOUR_PERIMETER:
                continue  # 周长不达标直接跳过

            # 3. 圆形度判断：核心优化（更严格的球体形状筛选）
            circularity = 4 * math.pi * area / (perimeter ** 2)
            # 圆形度阈值提高到0.65，减少非球体干扰
            if circularity < 0.65:
                continue  # 非圆形直接跳过

            # 4. 新增：轮廓外接矩形宽高比（排除扁长形）
            x, y, w, h = cv2.boundingRect(contour)
            aspect_ratio = min(w, h) / max(w, h)  # 宽高比（越接近1越圆）
            if aspect_ratio < 0.7:  # 过滤扁长形轮廓
                continue

            # 5. 计算质心（圆心）
            M = cv2.moments(contour)
            if M["m00"] <= 0:
                continue  # 避免除零错误
            cx_small = int(M["m10"] / M["m00"])
            cy_small = int(M["m01"] / M["m00"])

            # 映射回原图像尺寸（640x480）
            cx = int(cx_small * (640 / 320))
            cy = int(cy_small * (480 / 240))

            # 6. 评分机制优化：综合面积、圆形度、紧凑度
            score = area * circularity * aspect_ratio  # 新增紧凑度权重
            if score > best_score:
                best_score = score
                ball_center = (cx, cy)

        return ball_center

    # -------------------------- 测距与显示逻辑（未修改） --------------------------
    def is_near_center(self, x, y):
        dx = x - 320
        dy = y - 240
        return math.sqrt(dx*dx + dy*dy) < self.center_threshold

    def measure_distance(self, frame):
        frame_left = frame[:, :640]
        frame_right = frame[:, 640:]

        gray_left = cv2.cvtColor(frame_left, cv2.COLOR_BGR2GRAY)
        gray_right = cv2.cvtColor(frame_right, cv2.COLOR_BGR2GRAY)

        left_rect = rectify_image(gray_left, self.cam_params["left_map1"], self.cam_params["left_map2"])
        right_rect = rectify_image(gray_right, self.cam_params["right_map1"], self.cam_params["right_map2"])

        _, _, disp_color, threeD = compute_disparity_and_depth(
            left_rect, right_rect, self.stereo, self.right_matcher, 
            self.wls_filter, self.cam_params["Q"]
        )

        ball_center = self.detect_orange_ball(frame_left)
        if not ball_center:
            return None, disp_color

        cx, cy = ball_center
        if 0 <= cx < 640 and 0 <= cy < 480:
            depth = threeD[cy][cx]
            z_val = depth[2]
            if z_val <= 0:
                return None, disp_color
            distance = math.sqrt(depth[0]**2 + depth[1]** 2 + depth[2]**2) / 1000.0
            return distance, disp_color

        return None, disp_color

    def image_callback(self, msg):
        try:
            frame = self.bridge.imgmsg_to_cv2(msg, 'bgr8')
            frame_left = frame[:, :640].copy()  # 复制避免原图修改
        except Exception as e:
            self.get_logger().error(f"Image error: {e}")
            return

        ball_center = self.detect_orange_ball(frame_left)
        depth_color = None

        with self.coord_lock:
            if ball_center:
                self.target_x, self.target_y = ball_center
                if self.is_near_center(self.target_x, self.target_y):
                    self.ball_distance, depth_color = self.measure_distance(frame)
                else:
                    self.ball_distance = None
            else:
                self.target_x, self.target_y = 320, 240
                self.ball_distance = None

        # -------------------------- 显示优化 --------------------------
        if ball_center:
            # 绘制小球相关标记（原逻辑保留）
            cv2.line(frame_left, (self.target_x, self.target_y), (320, 240), 
                     (0, 255, 255), 2)  # 黄色连线
            cv2.circle(frame_left, (self.target_x, self.target_y), 10, (0, 255, 0), 2)
            cv2.circle(frame_left, (self.target_x, self.target_y), 3, (0, 0, 255), -1)

            # 显示坐标和距离（原逻辑保留）
            cv2.putText(frame_left, f"Center: ({self.target_x}, {self.target_y})", 
                       (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
            if self.ball_distance:
                cv2.putText(frame_left, f"Distance: {self.ball_distance:.2f}m", 
                           (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        else:
            cv2.putText(frame_left, "NO TARGET", (10, 30),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)

        # 绘制屏幕中心标记（原逻辑保留）
        cv2.circle(frame_left, (320, 240), 5, (255, 0, 0), -1)
        cv2.line(frame_left, (310, 240), (330, 240), (255, 0, 0), 2)
        cv2.line(frame_left, (320, 230), (320, 250), (255, 0, 0), 2)

        # -------------------------- 深度图显示（新增距离文本） --------------------------
        if self.display_enabled:
            cv2.imshow("Ball Detection", frame_left)
            if depth_color is not None:
                # 若距离有效，在深度图上叠加文本
                if self.ball_distance is not None:
                    cv2.putText(
                        depth_color, 
                        f"Distance: {self.ball_distance:.2f}m", 
                        (10, 30),  # 文本位置（左上角）
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7,  # 字体样式
                        (255, 255, 255), 2  # 白色字体，加粗
                    )
                cv2.imshow("Depth Map", depth_color)
            cv2.waitKey(1)

    # -------------------------- 串口发送逻辑（未修改） --------------------------
    def send_coordinates_loop(self):
        while rclpy.ok():
            with self.coord_lock:
                x, y = self.target_x, self.target_y
            self.send_coordinate(x, y)
            time.sleep(SEND_INTERVAL)

    def send_coordinate(self, x, y):
        if not self.ser or not self.ser.is_open:
            return

        x_clamped = np.clip(int(x), 0, 65535)
        y_clamped = np.clip(int(y), 0, 65535)

        x_low = x_clamped & 0xFF
        x_high = (x_clamped >> 8) & 0xFF
        y_low = y_clamped & 0xFF
        y_high = (y_clamped >> 8) & 0xFF

        frame = bytearray([FRAME_HEADER, FRAME_LENGTH, x_low, x_high, y_low, y_high, FRAME_TAIL])
        try:
            self.ser.write(frame)
            self.get_logger().debug(f"Sent: ({x_clamped}, {y_clamped})")
        except serial.SerialException as e:
            self.get_logger().error(f"Serial send failed: {e}")

    def destroy_node(self):
        if self.display_enabled:
            cv2.destroyAllWindows()
        if self.ser and self.ser.is_open:
            self.ser.close()
        super().destroy_node()


def main(args=None):
    rclpy.init(args=args)
    node = BallDetectorNode()
    executor = MultiThreadedExecutor(num_threads=2)
    executor.add_node(node)
    try:
        executor.spin()
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
